from __future__ import annotations

from typing import Iterable, Callable, TYPE_CHECKING

if TYPE_CHECKING:
    import pandas as pd

class DataSet:
    def __init__(self) -> None:
        self.dataset: dict[int, dict[str, pd.DataFrame]] = {}
        self.depth_map = {}
        self.name_stage_map: dict[str, int] = {}

    def regist(self, data: pd.DataFrame, name: str, stage: int) -> None:
        stage_datas = self.dataset.setdefault(stage, {})
        if name in stage_datas:
            print(f'Warning! {name} is already exist in stage {stage}, overwrite')
        depth = len(name.split('/'))
        depth_map = self.depth_map.setdefault(depth, {})
        stage_datas[name] = data
        depth_map[name] = data
        self.name_stage_map[name] = stage

    def get_from_stage(self, stage: int) -> Iterable[DataWrapper]:
        if stage in self.dataset:
            for name, data in self.dataset[stage].items():
                yield DataWrapper(data, self, name)
        else:
            raise Exception(f'{stage} not exist')

    def get_from_stages(self, stages: list[int]) -> Iterable[DataWrapper]:
        for stage in stages:
            if stage in self.dataset:
                for name, data in self.dataset[stage].items():
                    yield DataWrapper(data, self, name)
            else:
                raise Exception(f'{stage} not exist')

    def get_data_from_name(self, name: str) -> DataWrapper:
        return self.dataset[self.name_stage_map[name]][name]

    def stage_mapping(self, stage_map: dict[int, int], name: str, fn: Callable, fn_params: dict) -> None:
        for source, target in stage_map.items():
            for datawrapper in self.get_from_stage(source):
                new_data = fn(datawrapper.data, **fn_params)
                datawrapper.append(new_data, name, target)

    def stage_size(self, stage: int) -> int:
        if stage in self.dataset:
            return len(self.dataset[stage])
        else:
            return 0

class DataWrapper:
    def __init__(self, data: pd.DataFrame, dataset: DataSet, name: str) -> None:
        self.data = data
        self.dataset = dataset
        self.name = name
        self.path = self.name.split('/')
        self.depth = len(self.path)
        self.stage = self.dataset.name_stage_map[self.name]

    def append(self, data: pd.DataFrame, name: str, stage: int) -> None:
        new_name = f'{self.name}/{name}'
        self.dataset.regist(data, new_name, stage)